package com.lagou.gateway.standard;

import com.lagou.gateway.constants.GatewayKnownHeader;
import com.lagou.gateway.standard.handler.ResponseHandlerFunction;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

public class GatewayResponseDecorator extends ServerHttpResponseDecorator {

    private ResponseHandlerFunction<DataBuffer> jsonBodyHandler;
    private ResponseHandlerFunction<DataBuffer> errorHandler;
    private ResponseHandlerFunction<DataBuffer> normalHandler;

    public GatewayResponseDecorator(
            ResponseHandlerFunction jsonBodyHandler,
            ResponseHandlerFunction errorHandler,
            ResponseHandlerFunction normalHandler,
            ServerHttpResponse delegate) {
        super(delegate);
        this.jsonBodyHandler = jsonBodyHandler;
        this.errorHandler = errorHandler;
        this.normalHandler = normalHandler;
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        if (getDelegate().getHeaders().containsKey(GatewayKnownHeader.GATEWAY_TOKEN)
                && getDelegate().getStatusCode().isError()) {
            return errorHandler.apply(getDelegate(), body);
        } else if (getDelegate().getHeaders().containsKey(GatewayKnownHeader.GATEWAY_TOKEN)
                && MediaType.APPLICATION_JSON_UTF8.equals(getDelegate().getHeaders().getContentType())
                ) {
            return jsonBodyHandler.apply(getDelegate(), body);
        } else {
            return normalHandler.apply(getDelegate(), body);
        }
    }

    @Override
    public Mono<Void> writeAndFlushWith(
            Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return writeWith(Flux.from(body).flatMapSequential(p -> p));
    }
}
